#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UCEIS数据集加载器
支持从文件夹结构自动加载图像和标签
"""

import os
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import torchvision.transforms as transforms
from pathlib import Path
from typing import Dict, Any, List, Tuple, Optional
import glob
import random

# BERT Tokenizer
from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')


class UCEISDataset(Dataset):
    """
    UCEIS评分数据集 - 反数据泄露版本
    强制使用文件列表，彻底避免数据泄露
    """

    def __init__(self, file_list: List[str], transform: Optional[transforms.Compose] = None, mode: str = 'train'):
        """
        初始化数据集（仅支持文件列表模式）

        Args:
            file_list: 包含图像文件绝对路径的列表（必须）
            transform: 图像变换
            mode: 'train' 或 'val'
        """
        if not isinstance(file_list, list):
            raise TypeError("UCEISDataset现在只接受文件列表，不再支持目录路径。请使用split_dataset_by_files函数。")

        if len(file_list) == 0:
            raise ValueError("文件列表不能为空！")

        self.file_list = file_list
        self.transform = transform
        self.mode = mode
        self.samples = []

        # 数据泄露检查
        self._validate_file_list()

        # 从文件列表构建样本
        self._build_samples_from_file_list()

        print(f"[反泄露] 数据集初始化完成: {len(self.samples)} 个样本")
        print(f"[反泄露] 类别分布: {self._get_class_distribution()}")

    def _validate_file_list(self):
        """验证文件列表的完整性和合法性"""
        print(f"[反泄露检查] 验证文件列表完整性...")

        # 检查文件存在性
        missing_files = []
        for file_path in self.file_list:
            if not Path(file_path).exists():
                missing_files.append(file_path)

        if missing_files:
            raise FileNotFoundError(f"发现{len(missing_files)}个缺失文件，数据可能被污染：{missing_files[:5]}...")

        # 检查重复文件
        if len(self.file_list) != len(set(self.file_list)):
            print("[警告] 发现重复文件路径，正在去重...")
            self.file_list = list(set(self.file_list))

        print(f"[反泄露检查] 文件列表验证通过：{len(self.file_list)}个唯一文件")

    def _build_samples_from_file_list(self):
        """从文件列表构建样本数据"""
        for img_path_str in self.file_list:
            img_path = Path(img_path_str)

            # 从文件路径推断分数
            score = self._extract_score_from_path(img_path)
            if score is None:
                continue  # 跳过无法识别分数的文件

            # 标签转换为0-7索引（1-8分 -> 0-7索引）
            label = score - 1

            sample = {
                'image_path': str(img_path),
                'score': score,
                'label': label,  # 0-7 for classification
                'clinical_text': f"Patient with UCEIS score {score}"  # 默认文本描述
            }

            self.samples.append(sample)

    def _extract_score_from_path(self, img_path: Path) -> Optional[int]:
        """从文件路径中提取UCEIS分数（增强版）"""
        path_str = str(img_path)

        # 详细调试信息
        if len(self.samples) < 5:  # 只打印前几个样本的调试信息
            print(f"[调试] 解析路径: {path_str}")

        # 检查文件是否有效（必须有文件名）
        if img_path.stem == '' or img_path.suffix != '.bmp':
            print(f"[警告] 无效的文件路径: {path_str}")
            return None

        # 按路径深度优先匹配（更精确的匹配）
        path_parts = path_str.replace('\\', '/').split('/')

        # 从路径的最深层开始匹配
        for i in range(len(path_parts)-1, -1, -1):
            part = path_parts[i]
            for score in range(1, 9):
                # 精确匹配文件夹名称 "X分"
                if part == f"{score}分":
                    if len(self.samples) < 5:
                        print(f"[调试] 精确匹配到分数: {score} (从文件夹: {part})")
                    return score

        # 备用方案：如果精确匹配失败，使用包含匹配
        for score in range(1, 9):
            if f"{score}分" in path_str:
                if len(self.samples) < 5:
                    print(f"[调试] 包含匹配到分数: {score}")
                return score

        # 如果没找到，记录警告
        print(f"[警告] 无法从路径中提取分数: {path_str}")
        return None

    def _scan_data_directory(self):
        """扫描数据目录，构建样本列表"""

        # 定义文件夹到分数的映射
        folder_to_score = {
            '1分': 1, '2分': 2, '3分': 3,
            '4分': 4, '5分': 5, '6分': 6, '7分': 7, '8分': 8
        }

        # 遍历所有子目录
        for subdir in self.data_dir.iterdir():
            if subdir.is_dir():
                # 处理"1分-3分"和"4分-8分"文件夹
                if subdir.name in ['1分-3分', '4分-8分']:
                    for score_folder in subdir.iterdir():
                        if score_folder.is_dir() and score_folder.name in folder_to_score:
                            score = folder_to_score[score_folder.name]
                            self._collect_images_from_folder(score_folder, score)

        # 如果没有找到样本，尝试直接在根目录下查找分数文件夹
        if not self.samples:
            for folder_name, score in folder_to_score.items():
                folder_path = self.data_dir / folder_name
                if folder_path.exists() and folder_path.is_dir():
                    self._collect_images_from_folder(folder_path, score)

    def _collect_images_from_folder(self, folder_path: Path, score: int):
        """从指定文件夹收集图像文件（支持递归搜索）"""
        # 首先尝试直接在文件夹中查找BMP文件
        bmp_files = list(folder_path.glob("*.bmp"))

        # 如果直接没有找到，递归搜索子文件夹
        if not bmp_files:
            bmp_files = list(folder_path.rglob("*.bmp"))

        for img_path in bmp_files:
            # 标签转换为0-7索引（1-8分 -> 0-7索引）
            label = score - 1

            sample = {
                'image_path': str(img_path),
                'score': score,
                'label': label,  # 0-7 for classification
                'clinical_text': f"Patient with UCEIS score {score}"  # 默认文本描述
            }

            self.samples.append(sample)

    def _get_class_distribution(self) -> Dict[int, int]:
        """获取类别分布统计"""
        distribution = {}
        for sample in self.samples:
            score = sample['score']
            distribution[score] = distribution.get(score, 0) + 1
        return distribution

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, str, torch.Tensor]:
        """
        获取单个样本

        Returns:
            image_tensor: 预处理后的图像张量
            text: 临床文本描述
            label: 分数标签（0-7）
        """
        sample = self.samples[idx]

        # 加载图像
        image = Image.open(sample['image_path']).convert('RGB')

        # 应用变换
        if self.transform:
            image = self.transform(image)

        # 获取文本和标签
        text = sample['clinical_text']
        label = torch.tensor(sample['label'], dtype=torch.long)

        return image, text, label


# 图像预处理变换
train_transform = transforms.Compose([
    # 数据增强
    transforms.RandomRotation(degrees=[-15, 15]),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.3),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),

    # TODO: Implement Custom Augmentation (Gaussian Noise / Specular Reflection Removal)

    # 调整尺寸和标准化
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


def custom_collate_fn(batch: List[Tuple], device: torch.device) -> Dict[str, Any]:
    """
    自定义批处理函数（Windows兼容版本）

    Args:
        batch: 批次样本列表 [(image, text, label), ...]
        device: 计算设备

    Returns:
        batch_dict: 包含批处理张量的字典
    """
    try:
        # 解包批次
        images, texts, labels = zip(*batch)

        # 处理图像和标签（保持在CPU上，让DataLoader的pin_memory处理传输）
        images = torch.stack(images)  # 保持在CPU
        labels = torch.stack(labels)  # 保持在CPU

        # 处理文本（BERT tokenization）- 添加错误处理
        try:
            text_batch = tokenizer(
                list(texts),
                padding='longest',
                truncation=True,
                max_length=512,
                return_tensors='pt'
            )
        except Exception as e:
            print(f"[Tokenizer错误] {e}")
            # 创建一个简单的dummy token
            text_batch = {
                'input_ids': torch.zeros(len(texts), 10, dtype=torch.long),
                'attention_mask': torch.ones(len(texts), 10, dtype=torch.long),
                'token_type_ids': torch.zeros(len(texts), 10, dtype=torch.long)
            }

        # 文本张量也保持在CPU上

        return {
            'images': images,
            'text_inputs': text_batch,
            'labels': labels,
            'device': device  # 传递设备信息，稍后移动
        }
    except Exception as e:
        print(f"❌ Collate函数错误: {e}")
        # 返回一个安全的默认batch
        batch_size = len(batch)
        return {
            'images': torch.zeros(batch_size, 3, 224, 224),
            'text_inputs': {
                'input_ids': torch.zeros(batch_size, 10, dtype=torch.long),
                'attention_mask': torch.ones(batch_size, 10, dtype=torch.long),
                'token_type_ids': torch.zeros(batch_size, 10, dtype=torch.long)
            },
            'labels': torch.zeros(batch_size, dtype=torch.long),
            'device': device
        }


def create_file_list_from_directory(data_dir: str) -> List[str]:
    """
    从数据目录创建所有图像文件的完整列表

    Args:
        data_dir: 数据根目录

    Returns:
        包含所有.bmp文件绝对路径的列表
    """
    data_path = Path(data_dir)
    all_bmp_files = []

    # 递归搜索所有.bmp文件
    for bmp_file in data_path.rglob("*.bmp"):
        all_bmp_files.append(str(bmp_file))

    return all_bmp_files


def split_dataset_by_files(data_dir: str, train_ratio: float = 0.9, random_seed: int = 42) -> Tuple[List[str], List[str]]:
    """
    在文件级别分割数据集，避免数据泄露（基于文件内容哈希的增强版）

    Args:
        data_dir: 数据根目录
        train_ratio: 训练集比例（推荐0.9）
        random_seed: 随机种子

    Returns:
        train_file_list, val_file_list: 分割后的训练和验证文件列表
    """
    import hashlib

    random.seed(random_seed)

    # 创建主文件列表
    all_bmp_files = create_file_list_from_directory(data_dir)

    if not all_bmp_files:
        raise FileNotFoundError(f"在 {data_dir} 中未找到任何 .bmp 文件。")

    print(f"[反泄露] 正在准备主文件列表...")
    print(f"[反泄露] 找到总文件数: {len(all_bmp_files)}")

    # 基于文件名和分数的唯一性去重（防止同名文件不同分数的数据泄露）
    print(f"[反泄露] 正在进行严格文件去重检查...")
    unique_file_entries = {}
    name_score_duplicates = 0
    path_duplicates = 0

    for file_path in all_bmp_files:
        try:
            # 提取文件名和分数
            file_name = Path(file_path).name
            score = _extract_score_from_path(Path(file_path))

            if score is None:
                print(f"[反泄露] 警告：无法提取分数，跳过文件: {file_path}")
                continue

            # 创建基于文件名和分数的唯一键
            name_score_key = f"{file_name}_{score}"

            # 检查是否存在同名但不同分数的文件（严重数据泄露风险）
            existing_entry = None
            for key, entry in unique_file_entries.items():
                if entry['name'] == file_name and entry['score'] != score:
                    # 发现同名但不同分数的文件
                    name_score_duplicates += 1
                    print(f"[反泄露] 严重警告: 发现同名但不同分数的文件!")
                    print(f"  现有: {entry['score']}分 -> {entry['path']}")
                    print(f"  新增: {score}分 -> {file_path}")
                    print(f"  文件名: {file_name}")

                    # 保留第一个遇到的，跳过后续的
                    existing_entry = entry
                    break

            if existing_entry is not None:
                # 跳过这个重复的文件
                continue

            # 检查完全相同的文件路径
            if name_score_key in unique_file_entries:
                path_duplicates += 1
                print(f"[反泄露] 发现完全相同的文件: {file_path}")
                continue

            # 添加到唯一文件列表
            unique_file_entries[name_score_key] = {
                'name': file_name,
                'score': score,
                'path': file_path
            }

        except Exception as e:
            print(f"[反泄露] 警告：无法处理文件 {file_path}: {e}")

    if name_score_duplicates > 0:
        print(f"[反泄露] 移除了 {name_score_duplicates} 个同名但不同分数的重复文件（数据泄露风险）")
    if path_duplicates > 0:
        print(f"[反泄露] 移除了 {path_duplicates} 个完全重复的文件")

    # 转换为文件路径列表
    unique_file_list = [entry['path'] for entry in unique_file_entries.values()]

    print(f"[反泄露] 去重后剩余文件数: {len(unique_file_list)}")

    if len(unique_file_list) == 0:
        raise ValueError("所有文件都被去重移除了！请检查数据集质量。")

    # 保证可复现的随机打乱
    random.shuffle(unique_file_list)

    # 创建分割
    split_index = int(len(unique_file_list) * train_ratio)
    train_file_list = unique_file_list[:split_index]
    val_file_list = unique_file_list[split_index:]

    print(f"[反泄露] 数据分割完毕: {len(train_file_list)} 个训练样本, {len(val_file_list)} 个验证样本。")

    # 重叠检测（终极安全检查）
    overlap = set(train_file_list) & set(val_file_list)
    if overlap:
        raise RuntimeError(f"[严重错误] 发现 {len(overlap)} 个重叠文件！数据泄露风险极高：{list(overlap)[:5]}")

    # 同名文件重叠检测
    train_names = {Path(f).name for f in train_file_list}
    val_names = {Path(f).name for f in val_file_list}
    name_overlap = train_names & val_names

    if name_overlap:
        print(f"[反泄露] 警告: 发现 {len(name_overlap)} 个同名文件在训练集和验证集中")
        print(f"[反泄露] 这可能导致数据泄露，建议进一步检查数据集质量")

        # 检查这些同名文件的分数是否一致
        inconsistent_files = []
        for name in list(name_overlap)[:5]:  # 只检查前5个
            train_files = [f for f in train_file_list if Path(f).name == name]
            val_files = [f for f in val_file_list if Path(f).name == name]

            for train_file in train_files:
                for val_file in val_files:
                    train_score = _extract_score_from_path(Path(train_file))
                    val_score = _extract_score_from_path(Path(val_file))

                    if train_score != val_score:
                        inconsistent_files.append({
                            'name': name,
                            'train_file': train_file,
                            'val_file': val_file,
                            'train_score': train_score,
                            'val_score': val_score
                        })

        if inconsistent_files:
            print(f"[反泄露] 严重错误: 发现 {len(inconsistent_files)} 个同名但不同分数的文件!")
            for info in inconsistent_files[:3]:
                print(f"  {info['name']}: 训练集{info['train_score']}分 vs 验证集{info['val_score']}分")
            raise RuntimeError("[严重错误] 数据集存在同名但不同分数的文件，无法进行可靠训练!")

    if len(val_file_list) == 0:
        raise ValueError("验证集为空！请检查数据量是否太少或分割比例是否正确。")

    # 分割验证
    total_ratio = len(train_file_list) / len(unique_file_list)
    print(f"[反泄露] 训练集比例: {total_ratio:.3f}, 验证集比例: {1-total_ratio:.3f}")

    return train_file_list, val_file_list


def _extract_score_from_path(img_path: Path) -> Optional[int]:
    """从文件路径中提取UCEIS分数"""
    path_str = str(img_path)

    # 按路径深度优先匹配（更精确的匹配）
    path_parts = path_str.replace('\\', '/').split('/')

    # 从路径的最深层开始匹配
    for i in range(len(path_parts)-1, -1, -1):
        part = path_parts[i]
        for score in range(1, 9):
            # 精确匹配文件夹名称 "X分"
            if part == f"{score}分":
                return score

    # 备用方案：如果精确匹配失败，使用包含匹配
    for score in range(1, 9):
        if f"{score}分" in path_str:
            return score

    return None


def split_dataset(dataset: UCEISDataset, train_ratio: float = 0.8, random_seed: int = 42) -> Tuple[UCEISDataset, UCEISDataset]:
    """
    分割数据集为训练集和验证集（向后兼容版本）

    Args:
        dataset: 原始数据集
        train_ratio: 训练集比例
        random_seed: 随机种子

    Returns:
        train_dataset, val_dataset: 分割后的训练集和验证集
    """
    random.seed(random_seed)

    # 获取样本索引
    indices = list(range(len(dataset)))
    random.shuffle(indices)

    # 计算分割点
    train_size = int(len(dataset) * train_ratio)
    train_indices = indices[:train_size]
    val_indices = indices[train_size:]

    # 创建子数据集
    train_dataset = torch.utils.data.Subset(dataset, train_indices)
    val_dataset = torch.utils.data.Subset(dataset, val_indices)

    return train_dataset, val_dataset